
#include "xdns-server-state-impl.h"
#include "xdns-data.h"
#include "xdns-decoder.h"
#include "xdns-header-info.h"
#include "xdns-protocol.h"
#include "xdns-query.h"
#include <assert.h>
#include <string.h>
#include <limits>

using namespace std;

const unsigned short XDNS_MAX_REQUEST_ID = std::numeric_limits<unsigned short>::max();

XdnsServerStateImpl::XdnsServerStateImpl(unsigned short port,
                                         const struct in_addr* address)
: XdnsServerState(port, address),
  m_nextRequestId(XDNS_MAX_REQUEST_ID),
  m_udpSocket(-1),
  m_relatedIds(),
  m_pendingQueries()
{
}

XdnsServerStateImpl::~XdnsServerStateImpl()
{
    if (m_udpSocket != -1)
    {
        close(m_udpSocket);
    }

    map<long, XdnsData*>::iterator iter = m_pendingQueries.begin();
    map<long, XdnsData*>::iterator end = m_pendingQueries.end();
    for (; iter != end; ++iter)
    {
        XdnsData* data = iter->second;
        if (data != NULL)
        {
            delete data;
        }
    }
    m_pendingQueries.clear();
}

XdnsError
XdnsServerStateImpl::SendQuery(long queryId,
                               const XdnsQuery& query)
{
    // Retrieve a valid request ID.
    unsigned short newRequestId = 0;
    map<unsigned short, long>::iterator newIter = m_relatedIds.end();
    XdnsError error = GetValidRequestId(newRequestId, newIter);
    if (error == XE_NoEnt)
    {
        /* Add the query to the pending queue. It should be processed
           when a query is responded and a request ID is freed.
        */
        return AddToPendingQueue(queryId, query);
    }
    else if (error != XE_Ok)
    {
        return error;
    }

    // Generate the data of the DNS query.
    XdnsData* data = NULL;
    error = XdnsData::Init(&data);
    if (error != XE_Ok)
    {
        return error;
    }
    assert(data != NULL);
    error = query.ComposeQueryData(newRequestId, *data);
    if (error != XE_Ok)
    {
        XdnsData::Release(data);
        return error;
    }

    /* Check if the data size is greater than the default size of
       one UDP packet. If yes, the data should be sent by TCP.
    */
    size_t udpOffset = query.GetUdpStartOffset();
    size_t dataSize = data->GetDataSize();
    assert(udpOffset < dataSize);
    if ((dataSize - udpOffset) > XDNS_UDP_PACKET_SIZE)
    {
        // TODO: TCP is not supported currently.
        XdnsData::Release(data);
        return XE_NoSupport;
    }

    error = SendDataViaUdp(*data, udpOffset);
    XdnsData::Release(data);
    if (error != XE_Ok)
    {
        return error;
    }

    m_relatedIds.insert(newIter, make_pair(newRequestId, queryId));
    return XE_Ok;
}

void
XdnsServerStateImpl::CancelQuery(long queryId)
{
    /* NOTE: Only the pending queries could be canceled. The
             requesting queries should not be canceled. Because
             the request ID needs to be released until receiving
             the response. The requesting queries will be ignored
             in the XdnsHandler.
    */
    map<long, XdnsData*>::iterator iter = m_pendingQueries.find(queryId);
    if (iter != m_pendingQueries.end())
    {
        XdnsData* data = iter->second;
        if (data != NULL)
        {
            delete data;
        }

        m_pendingQueries.erase(iter);
    }
}

XdnsError
XdnsServerStateImpl::GetProcessingFds(int* nfds,
                                      fd_set* readfds,
                                      fd_set* writefds)
{
    if ((NULL == nfds) ||
        (NULL == readfds) ||
        (NULL == writefds))
    {
        return XE_Inval;
    }

    if ((!(m_relatedIds.empty())) &&
        (m_udpSocket != -1))
    {
        FD_SET(m_udpSocket, readfds);
        *nfds = m_udpSocket + 1;
        return XE_Ok;
    }

    return XE_NoEnt;
}

XdnsError
XdnsServerStateImpl::Process(fd_set* readfds,
                             fd_set* writefds)
{
    if (m_relatedIds.empty())
    {
        return XE_Ok;
    }

    return RecvDataViaUdp(readfds);
}

/*! Retrieve a new valid request ID.

    @param requestId On return a valid request ID.
    @param iter On return an iterator used to insert the related IDs.
    @return XE_Ok if success,
            XE_NoEnt if the request ID is collided with IDs in map.
*/
XdnsError
XdnsServerStateImpl::GetValidRequestId(unsigned short& requestId,
                                       map<unsigned short, long>::iterator& iter)
{
    long newRequestId = NextRequestId();
    map<unsigned short, long>::iterator newIter = m_relatedIds.lower_bound(newRequestId);
    if ((newIter != m_relatedIds.end()) &&
        (newIter->first == newRequestId))
    {
        /* The new request ID is present, reset the next request ID and retrieve
           another one.
        */
        m_nextRequestId = XDNS_MAX_REQUEST_ID;
        newRequestId = NextRequestId();
        newIter = m_relatedIds.lower_bound(newRequestId);
        if ((newIter != m_relatedIds.end()) &&
           (newIter->first == newRequestId))
        {
            return XE_NoEnt;
        }
    }

    requestId = newRequestId;
    iter = newIter;
    return XE_Ok;
}

unsigned short
XdnsServerStateImpl::NextRequestId()
{
    return (m_nextRequestId--);
}

XdnsError
XdnsServerStateImpl::OpenUdpSocket()
{
    if (m_udpSocket != -1)
    {
        // The socket has already been opened.
        return XE_Ok;
    }

    int udpSocket = socket(AF_INET, SOCK_DGRAM, 0);
    if (-1 == udpSocket)
    {
        return XE_Network;
    }

    int reuse = 1;
    if(-1 == setsockopt(udpSocket, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(int)))
    {
        close(udpSocket);
        return XE_Network;
    }

    struct sockaddr_in addrToConn = {0};
    memset(&addrToConn, 0, sizeof(addrToConn));
    addrToConn.sin_family = AF_INET;
    memcpy(&(addrToConn.sin_addr), &m_address, sizeof(addrToConn.sin_addr));
    addrToConn.sin_port = m_port;

    if (-1 == connect(udpSocket, (struct sockaddr*) (&addrToConn), sizeof(addrToConn)))
    {
        close(udpSocket);
        return XE_Network;
    }

    m_udpSocket = udpSocket;
    return XE_Ok;
}

XdnsError
XdnsServerStateImpl::SendDataViaUdp(const XdnsData& data,
                                    size_t startOffset)
{
    size_t dataSize = data.GetDataSize();
    if (dataSize <= startOffset)
    {
        return XE_Inval;
    }

    size_t udpDataSize = dataSize - startOffset;
    if ((dataSize - startOffset) > XDNS_UDP_PACKET_SIZE)
    {
        return XE_Inval;
    }

    unsigned char* dataBuf = new (nothrow) unsigned char[udpDataSize];
    if (NULL == dataBuf)
    {
        return XE_NoMem;
    }
    if (data.GetData(dataBuf, startOffset, udpDataSize) != udpDataSize)
    {
        delete [] dataBuf;
        return XE_Unexpected;
    }

    XdnsError error = OpenUdpSocket();
    if (error != XE_Ok)
    {
        delete [] dataBuf;
        return error;
    }

    ssize_t sentSize = send(m_udpSocket, dataBuf, udpDataSize, 0);
    delete [] dataBuf;
    if (sentSize != static_cast<ssize_t>(udpDataSize))
    {
        return XE_Network;
    }

    return XE_Ok;
}

XdnsError
XdnsServerStateImpl::RecvDataViaUdp(fd_set* readfds)
{
    if (NULL == readfds)
    {
        return XE_Inval;
    }

    if ((-1 == m_udpSocket) ||
        (!FD_ISSET(m_udpSocket, readfds)))
    {
        return XE_Ok;
    }

    unsigned char dataBuf[XDNS_UDP_PACKET_SIZE] = {0};
    ssize_t recvSize = recv(m_udpSocket, dataBuf, sizeof(dataBuf), 0);
    if (recvSize < XDNS_HEADER_SIZE)
    {
        return XE_BadResponse;
    }

    XdnsData* data = NULL;
    XdnsError error = XdnsData::Init(&data);
    if (error != XE_Ok)
    {
        return error;
    }
    assert(data != NULL);

    error = data->AppendData(dataBuf, static_cast<size_t>(recvSize));
    if (error != XE_Ok)
    {
        XdnsData::Release(data);
        return error;
    }

    error = ProcessReply(*data, 0);
    XdnsData::Release(data);
    return error;
}

/*! Add the query to the pending queue.

    This function adds the query to the pending queue. It is only called
    when no request IDs are available. the pending queries could be
    processed when a query is responded and a request ID is freed.

    @param queryId ID to identify the query.
    @param query Query to add.
    @return XE_Ok if success.
*/
XdnsError
XdnsServerStateImpl::AddToPendingQueue(long queryId,
                                       const XdnsQuery& query)
{
    // Generate the data of the DNS query with a temporary request ID.
    XdnsData* data = NULL;
    XdnsError error = XdnsData::Init(&data);
    if (error != XE_Ok)
    {
        return error;
    }
    assert(data != NULL);
    error = query.ComposeQueryData(XDNS_MAX_REQUEST_ID, *data);
    if (error != XE_Ok)
    {
        XdnsData::Release(data);
        return error;
    }

    if (!(m_pendingQueries.insert(make_pair(queryId, data)).second))
    {
        XdnsData::Release(data);
        return XE_Unexpected;
    }

    return XE_Ok;
}

XdnsError
XdnsServerStateImpl::ProcessReply(const XdnsData& data,
                                  size_t startOffset)
{
    XdnsHeaderInfo header;
    XdnsError error = XdnsDecoder::DecodeHeader(data, startOffset, header);
    if (error != XE_Ok)
    {
        return error;
    }

    map<unsigned short, long>::iterator iter = m_relatedIds.find(header.m_id);
    if (iter == m_relatedIds.end())
    {
        // The query is canceled.
        return XE_Ok;
    }

    if (header.m_tc)
    {
        error = XE_UdpTruncated;
    }
    else
    {
        switch (header.m_rcode)
        {
        case XDNS_RCODE_NO_ERROR:
        {
            break;
        }

        case XDNS_RCODE_FORMAT_ERROR:
        {
            error = XE_RcodeFormatError;
            break;
        }

        case XDNS_RCODE_SERVER_FAILURE:
        {
            error = XE_RcodeServerFailure;
            break;
        }

        case XDNS_RCODE_NAME_ERROR:
        {
            error = XE_RcodeNameError;
            break;
        }

        case XDNS_RCODE_NOT_IMPLEMENTED:
        {
            error = XE_RcodeNotImplemented;
            break;
        }

        case XDNS_RCODE_REFUSED:
        {
            error = XE_RcodeRefused;
            break;
        }

        default:
        {
            error = XE_NoSupport;
            break;
        }
        }
    }

    NotifyListeners(iter->second, error, &data, startOffset, &header);
    m_relatedIds.erase(iter);

    return XE_Ok;
}

